#include "../include/ASTNode.h"
#include "Module.h"
#include "Function.h"
#include "BasicBlock.h"
#include "IRbuilder.h"
#include "Type.h"
#include "Constant.h"
#include "Instruction.h"
#include "GlobalVariable.h"
#include <iostream>
#include <map>

// 全局变量用于存储当前的模块和函数
static Module* current_module = nullptr;
static Function* current_function = nullptr;

// 符号表：变量名 -> Value*
static std::map<std::string, Value*> symbol_table;

// 常量表：常量名 -> Value*
static std::map<std::string, Value*> const_table;

// ProgramNode实现
Value* ProgramNode::generate_ir(IRBuilder* builder) {
    // 获取模块（从builder）
    current_module = builder->get_module();
    
    // 遍历子节点
    for (auto& child : children) {
        child->generate_ir(builder);
    }
    
    return nullptr;
}

// CompUnitNode实现
Value* CompUnitNode::generate_ir(IRBuilder* builder) {
    // 遍历子节点（通常是函数定义和全局声明）
    for (auto& child : children) {
        child->generate_ir(builder);
    }
    return nullptr;
}

// FuncDefNode实现  
Value* FuncDefNode::generate_ir(IRBuilder* builder) {
    // 获取函数名
    std::string func_name = value;
    
    // 获取返回类型（从第一个子节点BasicType获取）
    Type* return_type = nullptr;
    if (!children.empty() && children[0]->node_type == "BasicType") {
        BasicTypeNode* type_node = static_cast<BasicTypeNode*>(children[0].get());
        return_type = type_node->get_type(current_module);
    }
    
    if (!return_type) {
        return_type = current_module->get_int32_type(); // 默认int类型
    }
    
    // 解析函数参数
    std::vector<Type*> param_types;
    std::vector<std::string> param_names;
    
    for (size_t i = 1; i < children.size(); ++i) {
        if (children[i]->node_type == "CompUnit") {
            FuncFParamsNode* params_node = static_cast<FuncFParamsNode*>(children[i].get());
            for (auto& param : params_node->children) {
                if (param->node_type == "FuncFParam") {
                    FuncFParamNode* param_node = static_cast<FuncFParamNode*>(param.get());
                    param_names.push_back(param_node->value);
                    if (!param_node->children.empty() && param_node->children[0]->node_type == "BasicType") {
                        BasicTypeNode* param_type_node = static_cast<BasicTypeNode*>(param_node->children[0].get());
                        param_types.push_back(param_type_node->get_type(current_module));
                    } else {
                        param_types.push_back(current_module->get_int32_type());
                    }
                }
            }
            break;
        }
    }
    
    // 创建函数类型
    auto func_type = FunctionType::get(return_type, param_types);
    
    // 创建函数
    current_function = Function::create(func_type, func_name, current_module);
    
    // 创建基本块
    auto entry_bb = BasicBlock::create(current_module, "entry", current_function);
    
    // 设置IRBuilder
    builder->set_insert_point(entry_bb);
    builder->set_curFunc(current_function);
    
    // 为函数参数创建符号表条目
    auto& func_args = current_function->get_args();
    auto arg_it = func_args.begin();
    for (size_t i = 0; i < param_names.size() && arg_it != func_args.end(); ++i, ++arg_it) {
        symbol_table[param_names[i]] = *arg_it;
    }
    
    // 处理函数体（Block节点）
    for (size_t i = 1; i < children.size(); ++i) {
        if (children[i]->node_type == "Block") {
            children[i]->generate_ir(builder);
        }
    }
    
    return current_function;
}

// FuncFParamNode实现
Value* FuncFParamNode::generate_ir(IRBuilder* builder) {
    // 参数的IR生成在FuncDefNode中处理
    return nullptr;
}

// FuncFParamsNode实现
Value* FuncFParamsNode::generate_ir(IRBuilder* builder) {
    // 参数列表的IR生成在FuncDefNode中处理
    return nullptr;
}

// BasicTypeNode实现
Value* BasicTypeNode::generate_ir(IRBuilder* builder) {
    // 这个节点不生成IR，只提供类型信息
    return nullptr;
}

Type* BasicTypeNode::get_type(Module* module) {
    if (value == "int") {
        return module->get_int32_type();
    } else if (value == "float") {
        return module->get_float_type();
    } else if (value == "void") {
        return module->get_void_type();
    }
    return module->get_int32_type(); // 默认返回int类型
}

// ConstDeclNode实现
Value* ConstDeclNode::generate_ir(IRBuilder* builder) {
    // 遍历常量定义
    for (auto& child : children) {
        child->generate_ir(builder);
    }
    return nullptr;
}

// VarDeclNode实现
Value* VarDeclNode::generate_ir(IRBuilder* builder) {
    // 遍历变量定义
    for (auto& child : children) {
        child->generate_ir(builder);
    }
    return nullptr;
}

// ConstDefNode实现
Value* ConstDefNode::generate_ir(IRBuilder* builder) {
    std::string const_name = value;
    
    // 获取初始化值
    Value* init_value = nullptr;
    if (!children.empty()) {
        init_value = children[0]->generate_ir(builder);
    }
    
    if (init_value) {
        const_table[const_name] = init_value;
    }
    
    return init_value;
}

// VarDefNode实现
Value* VarDefNode::generate_ir(IRBuilder* builder) {
    std::string var_name = value;
    
    // 判断是否在函数内部（局部变量）还是全局变量
    if (current_function) {
        // 局部变量：创建 alloca 指令
        Type* var_type = current_module->get_int32_type(); // 默认int类型
        Value* alloca_inst = builder->create_alloca(var_type);
        
        // 添加到符号表
        symbol_table[var_name] = alloca_inst;
        
        // 如果有初始化值
        if (!children.empty()) {
            Value* init_value = children[0]->generate_ir(builder);
            if (init_value) {
                builder->create_store(init_value, alloca_inst);
            }
        }
        
        return alloca_inst;
    } else {
        // 全局变量：创建全局变量
        Type* var_type = current_module->get_int32_type();
        Value* init_value = nullptr;
        
        if (!children.empty()) {
            init_value = children[0]->generate_ir(builder);
        }
        
        if (!init_value) {
            init_value = ConstantInt::get(0, current_module);
        }
        
        GlobalVariable* global_var = GlobalVariable::create(
            var_name, current_module, var_type, false, 
            static_cast<Constant*>(init_value)
        );
        
        symbol_table[var_name] = global_var;
        return global_var;
    }
}

// InitValNode实现
Value* InitValNode::generate_ir(IRBuilder* builder) {
    // 遍历子节点生成初始化值
    if (!children.empty()) {
        return children[0]->generate_ir(builder);
    }
    return nullptr;
}

// ConstInitValNode实现
Value* ConstInitValNode::generate_ir(IRBuilder* builder) {
    // 遍历子节点生成常量初始化值
    if (!children.empty()) {
        return children[0]->generate_ir(builder);
    }
    return nullptr;
}

// BlockNode实现
Value* BlockNode::generate_ir(IRBuilder* builder) {
    // 遍历块中的所有语句
    for (auto& child : children) {
        child->generate_ir(builder);
    }
    return nullptr;
}

// AssignStmtNode实现
Value* AssignStmtNode::generate_ir(IRBuilder* builder) {
    if (children.size() >= 2) {
        // 第一个子节点是左值（标识符），第二个是右值（表达式）
        Value* rvalue = children[1]->generate_ir(builder);
        
        if (children[0]->node_type == "IdentExpr") {
            IdentExprNode* ident_node = static_cast<IdentExprNode*>(children[0].get());
            std::string var_name = ident_node->value;
            
            auto it = symbol_table.find(var_name);
            if (it != symbol_table.end()) {
                return builder->create_store(rvalue, it->second);
            }
        }
    }
    return nullptr;
}

// ExprStmtNode实现
Value* ExprStmtNode::generate_ir(IRBuilder* builder) {
    // 执行表达式（忽略返回值）
    if (!children.empty()) {
        return children[0]->generate_ir(builder);
    }
    return nullptr;
}

// ReturnStmtNode实现
Value* ReturnStmtNode::generate_ir(IRBuilder* builder) {
    Value* return_value = nullptr;
    
    // 如果有返回值表达式
    if (!children.empty()) {
        return_value = children[0]->generate_ir(builder);
    }
    
    // 生成return指令
    if (return_value) {
        return builder->create_ret(return_value);
    } else {
        return builder->create_void_ret();
    }
}

// IfStmtNode实现
Value* IfStmtNode::generate_ir(IRBuilder* builder) {
    if (children.size() < 2) return nullptr;
    
    // 条件表达式
    Value* cond_value = children[0]->generate_ir(builder);
    
    // 创建基本块
    BasicBlock* then_bb = BasicBlock::create(current_module, "if.then", current_function);
    BasicBlock* else_bb = nullptr;
    BasicBlock* end_bb = BasicBlock::create(current_module, "if.end", current_function);
    
    if (children.size() >= 3) {
        else_bb = BasicBlock::create(current_module, "if.else", current_function);
    }
    
    // 生成条件分支
    if (else_bb) {
        builder->create_cond_br(cond_value, then_bb, else_bb);
    } else {
        builder->create_cond_br(cond_value, then_bb, end_bb);
    }
    
    // 生成then块
    builder->set_insert_point(then_bb);
    children[1]->generate_ir(builder);
    builder->create_br(end_bb);
    
    // 生成else块（如果存在）
    if (else_bb && children.size() >= 3) {
        builder->set_insert_point(else_bb);
        children[2]->generate_ir(builder);
        builder->create_br(end_bb);
    }
    
    // 设置插入点到end块
    builder->set_insert_point(end_bb);
    
    return nullptr;
}

// BinaryExprNode实现
Value* BinaryExprNode::generate_ir(IRBuilder* builder) {
    if (children.size() < 2) return nullptr;
    
    Value* left = children[0]->generate_ir(builder);
    Value* right = children[1]->generate_ir(builder);
    
    if (!left || !right) return nullptr;
    
    // 根据操作符生成对应的指令
    if (op == "+") {
        return builder->create_iadd(left, right);
    } else if (op == "-") {
        return builder->create_isub(left, right);
    } else if (op == "*") {
        return builder->create_imul(left, right);
    } else if (op == "/") {
        return builder->create_isdiv(left, right);
    } else if (op == "%") {
        return builder->create_irem(left, right);
    } else if (op == "<") {
        return builder->create_icmp_lt(left, right);
    } else if (op == "<=") {
        return builder->create_icmp_le(left, right);
    } else if (op == ">") {
        return builder->create_icmp_gt(left, right);
    } else if (op == ">=") {
        return builder->create_icmp_ge(left, right);
    } else if (op == "==") {
        return builder->create_icmp_eq(left, right);
    } else if (op == "!=") {
        return builder->create_icmp_ne(left, right);
    } else if (op == "&&") {
        // 短路求值的简化实现
        return builder->create_iand(left, right);
    } else if (op == "||") {
        // 短路求值的简化实现
        return builder->create_ior(left, right);
    }
    
    return nullptr;
}

// UnaryExprNode实现
Value* UnaryExprNode::generate_ir(IRBuilder* builder) {
    if (children.empty()) return nullptr;
    
    Value* operand = children[0]->generate_ir(builder);
    if (!operand) return nullptr;
    
    if (op == "+") {
        return operand; // 正号不做任何操作
    } else if (op == "-") {
        Value* zero = ConstantInt::get(0, current_module);
        return builder->create_isub(zero, operand);
    } else if (op == "!") {
        // 使用逻辑非：!x = (x == 0)
        Value* zero = ConstantInt::get(0, current_module);
        return builder->create_icmp_eq(operand, zero);
    }
    
    return nullptr;
}

// CallExprNode实现
Value* CallExprNode::generate_ir(IRBuilder* builder) {
    std::string func_name = value;
    
    // 查找函数
    Function* callee = nullptr;
    auto functions = current_module->get_functions();
    for (auto func : functions) {
        if (func->get_name() == func_name) {
            callee = func;
            break;
        }
    }
    
    if (!callee) {
        std::cerr << "Error: Function " << func_name << " not found" << std::endl;
        return nullptr;
    }
    
    // 生成参数
    std::vector<Value*> args;
    for (auto& child : children) {
        Value* arg = child->generate_ir(builder);
        if (arg) {
            args.push_back(arg);
        }
    }
    
    return builder->create_call(callee, args);
}

// IdentExprNode实现
Value* IdentExprNode::generate_ir(IRBuilder* builder) {
    std::string var_name = value;
    
    // 首先查找常量表
    auto const_it = const_table.find(var_name);
    if (const_it != const_table.end()) {
        return const_it->second;
    }
    
    // 然后查找变量符号表
    auto it = symbol_table.find(var_name);
    if (it != symbol_table.end()) {
        Value* var = it->second;
        
        // 检查是否是全局变量或局部变量指针
        if (var->get_type()->is_pointer_type()) {
            // 生成load指令
            return builder->create_load(var);
        } else {
            // 直接返回值（用于函数参数等）
            return var;
        }
    }
    
    std::cerr << "Error: Variable " << var_name << " not found" << std::endl;
    return nullptr;
}

// LiteralExprNode实现
Value* LiteralExprNode::generate_ir(IRBuilder* builder) {
    if (literal_type == "int") {
        int int_value = std::stoi(value);
        return ConstantInt::get(int_value, current_module);
    } else if (literal_type == "float") {
        // 暂时不支持浮点常量，转换为整数
        int int_value = static_cast<int>(std::stof(value));
        return ConstantInt::get(int_value, current_module);
    }
    
    // 默认返回0
    return ConstantInt::get(0, current_module);
}

// ParenExprNode实现
Value* ParenExprNode::generate_ir(IRBuilder* builder) {
    // 括号表达式直接返回子表达式的值
    if (!children.empty()) {
        return children[0]->generate_ir(builder);
    }
    return nullptr;
} 